# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: inkyu

from abc import abstractmethod
from collections import namedtuple
import enum

from BTrees.OOBTree import OOBTree


class SizeUnit(enum.Enum):
  QTY = 1
  AMT = 2


PriceQty = namedtuple('PriceQty', ['price', 'qty'])
NotionalSize = namedtuple('NotionalSize', ['accumulated', 'entries', 'unit'])


class PriceMapIterator(object):
  def __init__(self, price_map):
    self._price_map = price_map
    self._iter = price_map.__iter__()

  def next(self):
    res = self._iter.next()
    print(res)


class PriceMap(object):
  def __init__(self):
    self.btree = OOBTree()
    self.clear_cache()

  @abstractmethod
  def price_to_key(self, price: float) -> float:
    pass

  @abstractmethod
  def key_to_price(self, key: float) -> float:
    pass

  def notify_updated(self, price, entry):
    self.clear_cache()

  def clear_cache(self):
    self._cache_min = None
    self._cache_notional_qty = (None, None)
    self._cache_notional_amt = (None, None)

  def clear(self):
    self.clear_cache()
    self.btree.clear()

  def get(self, price: float, default=None):
    return self.btree.get(self.price_to_key(price), default)

  def not_empty(self):
    return self.btree.__bool__()

  def __contains__(self, price: float):
    return self.price_to_key(price) in self.btree

  def __getitem__(self, price: float):
    return self.btree[self.price_to_key(price)]

  def __setitem__(self, price: float, entry):
    self.clear_cache()
    self.btree[self.price_to_key(price)] = entry

  def __delitem__(self, price: float):
    self.clear_cache()
    del self.btree[self.price_to_key(price)]

  def iteritems(self, *args, **kwargs):
    for key, value in self.btree.iteritems(*args, **kwargs):
      yield (self.key_to_price(key), value)

  def get_min(self):
    if self._cache_min is not None:
      return self._cache_min

    if self.not_empty():
      min_key = self.btree.minKey()
      value = self.btree[min_key]
      self._cache_min = value
      return value

    else:
      return None

  def get_entries(self, from_price=None, to_price=None, exclude_max=False):
    from_price = self.price_to_key(from_price)
    to_price = self.price_to_key(to_price)
    min_price = min(from_price, to_price)
    max_price = max(from_price, to_price)
    return self.btree.values(min=min_price, max=max_price, excludemin=False, excludemax=exclude_max)

  def get_notional_entries_by_qty(self, qty):
    if self._cache_notional_qty[0] == qty:
      return self._cache_notional_qty[1]

    accumulated_qty = 0
    result = []
    for key, value in self.iteritems():
      assert key == value.price
      cur_price = key
      cur_qty = value.qty
      assert cur_qty > 0
      remaining_qty = qty - accumulated_qty
      if cur_qty >= remaining_qty:
        accumulated_qty += remaining_qty
        result.append(PriceQty(cur_price, remaining_qty))
        break
      else:
        accumulated_qty += cur_qty
        result.append(PriceQty(cur_price, cur_qty))

    ret = NotionalSize(accumulated_qty, result, unit=SizeUnit.QTY)
    self._cache_notional_qty = (qty, ret)
    return ret

  def get_notional_entries_by_amt(self, amt):
    if self._cache_notional_amt[0] == amt:
      return self._cache_notional_amt[1]

    accumulated_amt = 0
    result = []
    for key, value in self.iteritems():
      assert key == value.price
      cur_price = key
      cur_qty = value.qty
      assert cur_qty > 0
      remaining_amt = amt - accumulated_amt
      if cur_qty * cur_price >= remaining_amt:
        accumulated_amt += remaining_amt
        result.append(PriceQty(cur_price, remaining_amt / cur_price))
        break
      else:
        accumulated_amt += cur_price * cur_qty
        result.append(PriceQty(cur_price, cur_qty))

    ret = NotionalSize(accumulated_amt, result, unit=SizeUnit.AMT)
    self._cache_notional_amt = (amt, ret)
    return ret


class BidPriceMap(PriceMap):
  def price_to_key(self, price):
    return -price

  def key_to_price(self, key):
    return -key


class AskPriceMap(PriceMap):
  def price_to_key(self, price):
    return price

  def key_to_price(self, key):
    return key
